Add aten._trilinear support to torch_lib core#2850
Conversation
|
@microsoft-github-policy-service agree |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2850 +/- ##
==========================================
+ Coverage 71.86% 72.65% +0.78%
==========================================
Files 239 259 +20
Lines 29139 31709 +2570
Branches 2875 2989 +114
==========================================
+ Hits 20942 23037 +2095
- Misses 7219 7659 +440
- Partials 978 1013 +35 ☔ View full report in Codecov by Harness. |
justinchuby
added
the
module: torchlib
There was a problem hiding this comment.
Pull request overview
Adds a torch_lib lowering for aten::_trilinear (via ONNX Einsum) and wires it into the torch-lib OpInfo-based test harness to exercise the new lowering.
Changes:
- Implement
aten::_trilinearlowering inonnxscript/function_libs/torch_lib/ops/core.pyby generating an einsum equation fromexpand*/sumdim. - Add OpInfo/sample-input coverage for
torch.ops.aten._trilinear.defaultintests/function_libs/torch_lib/extra_opinfo.py. - Register the new op test entry in
tests/function_libs/torch_lib/ops_test_data.py.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| tests/function_libs/torch_lib/ops_test_data.py | Registers _trilinear in the torch_lib op test list. |
| tests/function_libs/torch_lib/extra_opinfo.py | Adds OpInfo + sample inputs to drive _trilinear tests. |
| onnxscript/function_libs/torch_lib/ops/core.py | Implements aten::_trilinear lowering using ONNX Einsum. |
|
Can you consider the copilot comments as well? |
There was a problem hiding this comment.
Pull request overview
This PR adds TorchLib support for PyTorch’s internal aten::_trilinear op by lowering it to ONNX Einsum, and extends the TorchLib OpInfo test matrix to cover the new op.
Changes:
- Add an opset18 TorchLib lowering for
aten::_trilinearinonnxscript/function_libs/torch_lib/ops/core.py. - Add OpInfo/sample-input coverage for
ops.aten._trilinear.defaultand register it in TorchLib test data. - Add dimension-list validation helpers for building the Einsum equation.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| tests/function_libs/torch_lib/ops_test_data.py | Registers ops.aten._trilinear.default in TorchLib operator test data. |
| tests/function_libs/torch_lib/extra_opinfo.py | Adds OpInfo + sample inputs for _trilinear to exercise tracing/export and numerical checks. |
| onnxscript/function_libs/torch_lib/ops/core.py | Implements _trilinear lowering via op.Einsum and adds helpers to build/validate the equation. |
| class _FakeTensor: | ||
| __slots__ = ("shape",) | ||
|
|
||
| def __init__(self, shape: Sequence[int]): | ||
| self.shape = shape |
There was a problem hiding this comment.
You may use ir.value() to create these objects, instead of creating a fake tensor
justinchuby
left a comment
justinchuby
left a comment
There was a problem hiding this comment.
Thanks! A nit comment
|
@WineChord could you check the previous comment before we can merge this? |
Summary
aten::_trilinearincore.pyexpand*/sumdimarguments as Python tuples during ONNX tracing tests_trilineardimension lists before building the einsum equationTesting
Context: requested from the review on pytorch/pytorch#177223.